#  Copyright (c) 2017-2018 Uber Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Minimal example of how to read samples from a dataset generated by `generate_hello_world_dataset.py`
using pyspark"""

from __future__ import print_function

from pyspark.sql import SparkSession

from examples.hello_world.generate_hello_world_dataset import HelloWorldSchema
from petastorm.spark_utils import dataset_as_rdd


def pyspark_hello_world(dataset_url='file:///tmp/hello_world_dataset'):
    spark = SparkSession \
        .builder \
        .master('local[1]') \
        .getOrCreate()

    # dataset_as_rdd creates an rdd of named tuples.
    rdd = dataset_as_rdd(dataset_url, spark, [HelloWorldSchema.id, HelloWorldSchema.image1])
    print('An id in the dataset: ', rdd.first().id)

    # Create a dataframe object from a parquet file
    dataframe = spark.read.parquet(dataset_url)

    # Show a schema
    dataframe.printSchema()

    # Count all
    dataframe.count()

    # Show just some columns
    dataframe.select('id').show()

    # This is how you can use a standard SQL to query a dataset. Note that the data is not decoded in this case.
    number_of_rows = spark.sql(
        'SELECT count(id) '
        'from parquet.`{}` '.format(dataset_url)).collect()
    print('Number of rows in the dataset: {}'.format(number_of_rows[0][0]))


if __name__ == '__main__':
    pyspark_hello_world()
